from math import sin, cos, tan, pi, sqrt
import matplotlib.pyplot as plt

plt.rc('font', family='Times New Roman')
plt.rc("axes", unicode_minus=False)
plt.rc("figure", figsize=(10,5))
plt.rc("savefig", dpi=400)
plt.rc("text", usetex = True)
ticklabels_style = {
    "fontname": "Times New Roman",
}

def a2r(phi):
    return phi/180.*pi

def s_phi(phi):
    #phi = a2r(phi)
    if phi>2*pi:
        return s_phi(phi-2*pi)
    elif phi>a2r(320):
        return 0
    elif phi>a2r(257.5):
        return 120*(phi-a2r(320))**2/a2r(125)/a2r(125)
    elif phi>a2r(195):
        return 60-120*(phi-a2r(195))**2/a2r(125)/a2r(125)
    elif phi>a2r(115):
        return 60
    elif phi>0:
        return 30*(1-cos(pi*phi/a2r(115)))
    else:
        return s_phi(phi+2*pi)

def v_phi(phi):
    #phi = a2r(phi)
    if phi>2*pi:
        return v_phi(phi-2*pi)
    elif phi>a2r(320):
        return 0
    elif phi>a2r(257.5):
        return 240*(phi-a2r(320))/a2r(125)/a2r(125)
    elif phi>a2r(195):
        return -240*(phi-a2r(195))/a2r(125)/a2r(125)
    elif phi>a2r(115):
        return 0
    elif phi>0:
        return 30*pi*sin(pi*phi/a2r(115))/a2r(115)
    else:
        return v_phi(phi+2*pi)

def a_phi(phi):
    ##phi = a2r(phi)
    if phi>2*pi:
        return a_phi(phi-2*pi)
    elif phi>a2r(320):
        return 0
    elif phi>a2r(257.5):
        return 240/a2r(125)/a2r(125)
    elif phi>a2r(195):
        return -240/a2r(125)/a2r(125)
    elif phi>a2r(115):
        return 0
    elif phi>0:
        return 30*pi*pi*cos(pi*phi/a2r(115))/a2r(115)/a2r(115)
    else:
        return a_phi(phi+2*pi)

def sv_slope(phi):
    return v_phi(phi)/a_phi(phi)

def sovle(slope,start,end,e=0.0001):
    while abs(start-end)>e:
        mid = (start+end)/2
        if (sv_slope(mid)-slope)*(sv_slope(start)-slope)<0:
            start,end = start,mid
        elif (sv_slope(mid)-slope)*(sv_slope(end)-slope)<0:
            start,end = mid,end
        elif (sv_slope(mid)-slope)==0:
            break
        else:
            return None
    return mid
        

gs_kw = dict(width_ratios=[1.5, 1], height_ratios=[1,1,1])

fig,axd = plt.subplot_mosaic(
    [['s-v','s'],
     ['s-v','v'],
     ['s-v','a']], gridspec_kw=gs_kw
)

axd['s-v'].set_aspect('equal')
axd['s-v'].set_xlim(-80, 80)
axd['s-v'].set_ylim(-60, 80)
axd['s'].set_ylim(0, 80)
axd['v'].set_ylim(-80, 80)
axd['a'].set_ylim(-100, 100)
axd['s'].set_xlim(0, 360)
axd['v'].set_xlim(0, 360)
axd['a'].set_xlim(0, 360)

dphi = 0.5
length = int(360/dphi)
phi_list = [a2r(i*dphi) for i in range(length)]
degree_list = [i*dphi for i in range(length)]

s_line, = axd  ['s'].plot(degree_list, list(map(s_phi,phi_list)), 'r-', lw=1, ms=0)
v_line, = axd  ['v'].plot(degree_list, list(map(v_phi,phi_list)), 'g-', lw=1, ms=0)
a_line, = axd  ['a'].plot(degree_list, list(map(a_phi,phi_list)), 'b-', lw=1, ms=0)
sv_line, = axd['s-v'].plot(list(map(v_phi,phi_list)), list(map(s_phi,phi_list)), 'c-', lw=1, ms=0)
sv_yaxis, = axd['s-v'].plot([0,0], [-100,100], 'grey', lw=1, ms=0)
sv_xaxis, = axd['s-v'].plot([-100,100], [0,0], 'grey', lw=1, ms=0)

phi1 = sovle(tan(a2r(90-35)),a2r(0.1),a2r(114.9))
phi2 = sovle(tan(a2r(90+70)),a2r(258),a2r(319.9))

edge1, = axd['s-v'].plot(
    [-100,100], 
    [tan(a2r(55))*(-100-v_phi(phi1))+s_phi(phi1), tan(a2r(55))*(100-v_phi(phi1))+s_phi(phi1)], 
    'orange', lw=1, ms=0, label='edge1')
edge2, = axd['s-v'].plot(
    [-100,100], 
    [tan(a2r(160))*(-100-v_phi(phi2))+s_phi(phi2), tan(a2r(160))*(100-v_phi(phi2))+s_phi(phi2)], 
    'orange', lw=1, ms=0, label='edge2')
edge3, = axd['s-v'].plot(
    [0,100], 
    [0,tan(a2r(125))*100], 
    'orange', lw=1, ms=0, label='edge2')
axd['s-v'].scatter([15],[-25])

axd['s-v'].set_xlabel(r'$\frac{ds}{d\varphi}/mm \cdot rad^{-1}$')
axd['s-v'].set_ylabel('s/mm')
axd['s'].set_xlabel(r'$\varphi$/°')
axd['s'].set_ylabel('s/mm')
axd['v'].set_xlabel(r'$\varphi$/°')
axd['v'].set_ylabel(r'$\frac{ds}{d\varphi}/mm \cdot rad^{-1}$')
axd['a'].set_xlabel(r'$\varphi$/°')
axd['a'].set_ylabel(r'$\frac{d^{2}s}{d\varphi^{2}}/mm \cdot rad^{-2}$')

plt.tight_layout()
plt.show()
plt.savefig('task2-base_select.png')

# choose r0=20sqrt(2),e=20 s0=20
